import a_trade_calendar
from config.env import StockConfig
from datetime import date, datetime, timedelta
import time
import pandas as pd
import akshare as ak
from utils.log_util import logger
from retrying import retry

@retry(stop_max_attempt_number=3, wait_fixed=3000)
def fetch_akdata(symbol, start_date=StockConfig.sync_ak_startdate, period='daily'):
    """
    从akshare获取股票x(日、周、月)线数据
    """
    try:
        # 避免请求过于频繁
        time.sleep(3)
        now = datetime.now()
        hour = now.hour
        end_date = now.strftime('%Y%m%d') if hour > 15 else (now - timedelta(days=1)).strftime('%Y%m%d')
        
        df = ak.stock_zh_a_hist(symbol=symbol, period=period , start_date=start_date, end_date=end_date, adjust="qfq", timeout=30)
        
        if df.empty:
            return pd.DataFrame()            
        rename_data(df)
        df.loc[(df['amplitude'] > 100000) | (df['amplitude'] < -100000) , 'amplitude'] = 0
        return df
    except Exception as e:
        logger.exception(e)
        raise e
    
def rename_data(old_data:pd.DataFrame):
    try:
        if 'index' in old_data.columns:
            old_data.drop('index',axis=1, inplace=True)
        old_data.rename(columns={
            "日期": "trade_date",
            "股票代码": "symbol",
            "开盘": "open",
            "收盘": "close",
            "最高": "high",
            "最低": "low",
            "成交量": "vol",
            "成交额": "amount",
            "振幅": "amplitude",
            "涨跌幅": "pct_chg",
            "涨跌额": "chg",
            "换手率": "turn"
            },inplace=True
        )
    except Exception as e:
        raise e

def filter_ak_stock(stock_df:pd.DataFrame, code_name:str, name_name:str):
    '''
    过滤股票集合
    1.只保留主板股票
    2.过滤掉有Nan数据的记录
    '''
    stock_df = stock_df.dropna()
    main_board_prefixes = ['600', '601', '603', '605', '000', '001', '002', '003']
    stock_df = stock_df[stock_df[code_name].astype(
        str).str.startswith(tuple(main_board_prefixes))]
    stock_df = stock_df[~stock_df[name_name].str.contains(
        'ST|退|PT') & ~stock_df[code_name].str.startswith(('400', '420'))]
    return stock_df

def filter_ak_rt_stock(stock_df:pd.DataFrame):
    '''
    过滤实时股票集合
    过滤掉最新价为0, 成交量为0,成交额为0的股票
    '''
    stock_df = stock_df[(stock_df["close"] > 0) & (stock_df["vol"]) > 0 & (stock_df["amount"] > 0) ]
    return stock_df

def get_latest_trade_date() :
    latest_trade_date_str = a_trade_calendar.get_latest_trade_date()
    format = "%Y-%m-%d"
    latest_trade_date = datetime.strptime(latest_trade_date_str, format).date()

    today_str = datetime.now().strftime(format)   
    hour = datetime.now().hour 
    is_today_trade = a_trade_calendar.is_trade_date(today_str)
    if is_today_trade :
        if hour > 15:
            return latest_trade_date
        else :
            return latest_trade_date - timedelta(days=1)
    else:
        return  latest_trade_date

@retry(stop_max_attempt_number=3)
def fetch_ak_rt_data():
    """
    从akshare获取股票最新日线实时数据
    """
    try:        
        df = ak.stock_zh_a_spot_em()
        if df.empty:
            return pd.DataFrame()            
        rename_rt_data(df)
        df.loc[(df['amplitude'] > 100000) | (df['amplitude'] < -100000) , 'amplitude'] = 0
        return df
    except Exception as e:
        logger.exception(e)
        raise e
    
def rename_rt_data(old_data:pd.DataFrame):
    try:
        if 'index' in old_data.columns:
            old_data.drop('index',axis=1, inplace=True)
        old_data.rename(columns={
            "代码": "symbol",
            "名称": "name",
            "今开": "open",
            "最新价": "close",
            "昨收":"pre_close",
            "最高": "high",
            "最低": "low",
            "成交量": "vol",
            "成交额": "amount",
            "振幅": "amplitude",
            "涨跌幅": "pct_chg",
            "涨跌额": "chg",
            "换手率": "turn"
            },inplace=True
        )
    except Exception as e:
        raise e
    